#include <cstddef>
#include <cstdint>
#include <iostream>
#include <map>
#include <memory>
#include <utility>
#include <vector>

using TensorHandle = int64_t;
using CtxHandle = int64_t;
constexpr TensorHandle kInvalidTensorHandle = -1;

class AscendTensor {
public:
  explicit AscendTensor(int64_t dtype = 0, std::vector<int64_t> dims = {})
      : dtype(dtype), dims(std::move(dims)) {}
  ~AscendTensor() = default;
  std::vector<int64_t> dims;
  int64_t dtype = 0;

  static std::shared_ptr<AscendTensor> Assign(
      const std::shared_ptr<AscendTensor> &x,
      const std::shared_ptr<AscendTensor> &y) {
    std::cout << "[AST] Assign(" << x.get() << ", " << y.get() << ")"
              << std::endl;
    return x;
  }

  static std::shared_ptr<AscendTensor> Add(
      const std::shared_ptr<AscendTensor> &x,
      const std::shared_ptr<AscendTensor> &y) {
    std::cout << "[AST] Add(" << x.get() << ", " << y.get() << ")" << std::endl;
    return std::make_shared<AscendTensor>();
  }

  static std::shared_ptr<AscendTensor> Mul(
      const std::shared_ptr<AscendTensor> &x,
      const std::shared_ptr<AscendTensor> &y) {
    std::cout << "[AST] Mul(" << x.get() << ", " << y.get() << ")" << std::endl;
    return std::make_shared<AscendTensor>();
  }
};

class Context {
public:
  Context() = default;
  ~Context() {
    for (size_t i = 0; i < tensors.size(); ++i) {
      std::cout << "[AST] Destroy tensor id=" << i << " " << tensors[i].get() << std::endl;
    }
  };
  virtual TensorHandle Add(const std::shared_ptr<AscendTensor> &tensor) {
    if (tensor == nullptr) {
      return kInvalidTensorHandle;
    }
    auto handle = TensorHandle(tensors.size());
    tensors.emplace_back(tensor);
    std::cout << "[AST] Create tensor id=" << handle << " " << tensor.get() << std::endl;
    return handle;
  }

  virtual std::shared_ptr<AscendTensor> Get(TensorHandle handle) {
    if (handle < 0 || handle >= tensors.size()) {
      return nullptr;
    }
    return tensors[handle];
  }

private:
  std::vector<std::shared_ptr<AscendTensor>> tensors;
};

class NullContext : public Context {
public:
  TensorHandle Add(const std::shared_ptr<AscendTensor> &tensor) override {
    return kInvalidTensorHandle;
  }

  std::shared_ptr<AscendTensor> Get(TensorHandle handle) override {
    return nullptr;
  }
};

class ContextManager {
public:
  static ContextManager &GetInstance() {
    static ContextManager instance;
    return instance;
  }

  std::shared_ptr<Context> Get(CtxHandle handle) {
    auto it = contexts.find(handle);
    if (it == contexts.end()) {
      return nullptr;
    }
    return it->second;
  }

  ~ContextManager() {
    for (auto &pair : contexts) {
      std::cout << "[AST] Destroy context " << pair.first << std::endl;
    }
    contexts.clear();
  }

  CtxHandle CreateContext() {
    auto ctx = std::make_shared<Context>();
    auto handle = CtxHandle(ctx.get());
    contexts[handle] = ctx;
    std::cout << "[AST] Create context " << handle << std::endl;
    return handle;
  }

  void DestroyContext(CtxHandle handle) {
    std::cout << "[AST] Destroy context " << handle << std::endl;
    contexts.erase(handle);
  }

private:
  ContextManager() = default;
  std::map<CtxHandle, std::shared_ptr<Context>> contexts;
};

namespace {
std::shared_ptr<Context> Ctx(CtxHandle handle) {
  static auto null_context = std::make_shared<NullContext>();
  auto ctx = ContextManager::GetInstance().Get(handle);
  if (ctx == nullptr) {
    std::cout << "Context not found:" << handle << std::endl;
    return null_context;
  }
  return ctx;
}
} // namespace

extern "C" {
int64_t GetTensorDtype(CtxHandle ctx_handle, TensorHandle handle) {
  auto ascend_tensor = Ctx(ctx_handle)->Get(handle);
  if (ascend_tensor == nullptr) {
    return -1;
  }
  return int64_t(ascend_tensor->dtype);
}

int64_t GetTensorDimNum(CtxHandle ctx_handle, TensorHandle handle) {
  auto ascend_tensor = Ctx(ctx_handle)->Get(handle);
  if (ascend_tensor == nullptr) {
    std::cout << "GetTensorDimNum failed" << std::endl;
    return -1;
  }
  return int64_t(ascend_tensor->dims.size());
}

int64_t GetTensorDims(CtxHandle ctx_handle, TensorHandle handle, int64_t dim_num, int64_t *dims) {
  auto ascend_tensor = Ctx(ctx_handle)->Get(handle);
  if (ascend_tensor == nullptr) {
    return -1;
  }
  if (dim_num < int64_t(ascend_tensor->dims.size()) || dims == nullptr) {
    return -1;
  }
  for (int64_t i = 0; i < dim_num; ++i) {
    dims[i] = ascend_tensor->dims[i];
  }
  return int64_t(ascend_tensor->dims.size());
}

CtxHandle CreateContext() { return ContextManager::GetInstance().CreateContext(); }

void DestroyContext(CtxHandle ctx_handle) { ContextManager::GetInstance().DestroyContext(ctx_handle); }

TensorHandle CreateTensor(CtxHandle ctx, int64_t dtype, int64_t dim_num, const int64_t *dims) {
  if (dims == nullptr || dim_num < 0) {
    return kInvalidTensorHandle;
  }
  std::vector<int64_t> vec_dims;
  vec_dims.reserve(dim_num);
  for (int64_t i = 0; i < dim_num; ++i) {
    vec_dims.push_back(dims[i]);
  }
  return Ctx(ctx)->Add(std::make_shared<AscendTensor>(dtype, vec_dims));
}

TensorHandle Assign(CtxHandle ctx_handle, TensorHandle x, TensorHandle y) {
  auto ctx = Ctx(ctx_handle);
  AscendTensor::Assign(ctx->Get(x), ctx->Get(y));
  return x;
}

TensorHandle Add(CtxHandle ctx_handle, TensorHandle x, TensorHandle y) {
  auto ctx = Ctx(ctx_handle);
  return ctx->Add(AscendTensor::Add(ctx->Get(x), ctx->Get(y)));
}

TensorHandle Mul(CtxHandle ctx_handle, TensorHandle x, TensorHandle y) {
  auto ctx = Ctx(ctx_handle);
  return ctx->Add(AscendTensor::Mul(ctx->Get(x), ctx->Get(y)));
}
};